/**
  MOMD project, Anyang Normal University, IMP-CAS
  \file ETFMath.h
  \class ETFMath
  \brief Math class, to provide some general math methods. Note that this is a
  tool class, so it is defined to be a static class.
  \author SUN Yazhou, aisa.rabbit@163.com
  \date Created: 2020/07/09
  \date Last modified: 2020/10/04 by SUN Yazhou
  \copyright 2020 SUN Yazhou
  \copyright MOMD project, Anyang Normal University, IMP-CAS
*/

#include <cmath>
#include <iostream>
#include <cstring>
#include <algorithm>
#include <complex> //complex number operation
#include <utility>
#include <array>
#include <vector>
#include "ETFMath.h"
#include "ETFMsg.h"

using std::swap;
using std::max;
using std::cout;
using std::endl;
using std::array;
using std::vector;
using std::complex;
typedef complex<double> cdouble;

// #define DEBUG

// the z border of the effective magField (in mm)
const double ETFMath::kzMagIn = -555.5;
const double ETFMath::kzMagOut = 555.5;

// r_global = R.r_local
// angle0: yaw, angle1: pitch, angle2: roll, intrinsic rotation; (y-x'-z")
// R(yaw, pitch, roll) = Ry(yaw).Rx'(pitch).Rz"(roll)
void ETFMath::rotate(const double *pIn, double *pOut, const double *angIn){
	double s1 = sin(angIn[0]), s2 = sin(angIn[1]), s3 = sin(angIn[2]);
	double c1 = cos(angIn[0]), c2 = cos(angIn[1]), c3 = cos(angIn[2]);

	pOut[0] = (c1*c3+s1*s2*s3)* pIn[0] +(c3*s1*s2-c1*s3)* pIn[1] +c2*s1* pIn[2];
	pOut[1] = c2*s3*            pIn[0] +c2*c3*            pIn[1] -s2*    pIn[2];
	pOut[2] = (c1*s2*s3-c3*s1)* pIn[0] +(c1*c3*s2+s1*s3)* pIn[1] +c1*c2* pIn[2];
} // end member function rotate

// print array in an orderly fashion
void ETFMath::peek(int *v, int len){
	for(int i = 0; i < len; i++){
		cout << v[i] << '\t';
		if(i != 0 && i % 10 == 0) cout << endl;
	} // end for
} // end member function Inspect

// returns the value ln[Gamma(xx)] for xx > 0.
double ETFMath::gammaln(double xx){
  if(xx <= 0.) ETFMsg::Error("ETFMath", "gammaln: x > 0. is mandatory.");
  // internal arithmetic will be done in double precision, a nicety that you can
  // omit if five-figure accuracy is good enough
  static const double c[6] = {76.18009172947146,-86.50532032941677,
    24.01409824083091,-1.231739572450155,
    0.1208650973866179e-2,-0.5395239384953e-5};
  double tmp = xx + 5.5, x = xx; // (z+gamma+1/2), gamma = 5, N = 6
  tmp -= (x+0.5)*log(tmp);
  double ser = 1.000000000190015; // the additive series
  for(int i = 0; i < 6; i++) ser += c[i]/++x;
  return -tmp+log(2.5066282746310005*ser/xx); // sqrt(2pi)
} // end member function gammaln

// returns ln(n!)
double ETFMath::factln(int n){
  static double a[101]; // a static array is automatically initialized to zero

  if(n < 0) ETFMsg::Error("ETFMath", "factln: n is negative.");
  if(n <= 1) return 0.;
  if(n <= 100) return a[n] ? a[n] : (a[n]=gammaln(n+1.));
  return gammaln(n+1.); // out of range of table
} // end member function factln

/// \retval n!
int ETFMath::Factorial(int n){
  static int ntop = 4;
  static int a[33] = {1, 1, 2, 6, 24}; // fill in table only as required

  if(n < 0) ETFMsg::Error("ETFMath", "Factorial: n: %d is minus", n);
  // larger value than size of the table is required. Actually, this big a value
  // is going to overflow on many computers, but no harm in trying
  if(n > 32) return exp(gammaln(n+1.));
  while(ntop < n){
    ntop++;
    a[ntop] = a[ntop-1]*ntop;
  }
  return floor(0.5+a[n]); // clear off the roundoff error
} // end member function Factorial

// returns Binomial coefficients
int ETFMath::Binomial(int n, int m){
  if(m > n)
    ETFMsg::Error("ETFMath",
      "Binomial: n: %d is larger than m: %d!", n, m);

  // the floor function cleans up roundoff error for smaller values of n and m
  return floor(0.5+exp(factln(n)-factln(m)-factln(n-m)));
} // end member function Binomial

double ETFMath::Beta(double z, double w){
  if(z < 0. || w < 0.) ETFMsg::Error("ETFMath", "Beta: z or z is minus");
  return exp(gammaln(z)+gammaln(w)-gammaln(z+w));
} // end member function Beta

/// \retval n!!
int ETFMath::BiFactorial(int n){
  return n <= 1 ? 1 : n * BiFactorial(n-2);
} // end member function BiFactorial

/// the incomplete gamma function gammap, gammap(a,x)=P(a,x)=gamma(a,x)/Gamma(a)
double ETFMath::gammap(double a, double x){
  if(x < 0. || a <= 0.) ETFMsg::Error("ETFMath", "gammap: invalid arguements");

  if(x < (a+1.)) return gser(a, x); // series representation
  return 1. - gcf(a, x);
} // end member function gammap
/// the incomplete gamma function gammaq, gammaq(a,x)=Q(a,x)=1-P(a,x)
double ETFMath::gammaq(double a, double x){
  if(x < 0. || a <= 0.) ETFMsg::Error("ETFMath", "gammaq: invalid arguements");

  if(x < (a+1.)) return 1. - gser(a, x); // continued fraction representation
  return gcf(a, x);
} // end member function gammaq

// returns the incomplete gamma function P(a,x) evaluated by its series
// representation. Optionally returns ln[Gamma(a)] as gln.
double ETFMath::gser(double a, double x, double *gln){
  static const double ITMAX = 100; // maximum iteration times
  static const double EPS = 3.e-7; // fractional accuracy

  if(x < 0.) ETFMsg::Error("ETFMath", "gser: x < 0., x: %f", x);
  double ap = a, sum, item = sum = 1./a;
  for(int i = 0; i < ITMAX; i++){
    sum += (item *= x/++ap);
    if(fabs(item) < fabs(sum)*EPS)
      return sum * exp(-x + a*log(x) - (gln ? *gln=gammaln(a) : gammaln(a)));
  } // end for over i
  ETFMsg::Error("ETFMath", "gser: a too large, and IMAX too small.");
  return 0.; // never gets here
} // end member function gamser

// returns the incomplete gamma function Q(a,x) evaluated by its continued
// fraction representation. Optionally returns ln[Gamma(a)] as gln.
double ETFMath::gcf(double a, double x, double *gln){
  static const double ITMAX = 100; // maximum iteration times
  static const double EPS = 3.e-7; // factional accuracy
  // number near the smallest representable floating-point number
  static const double FPMIN = 1.e-30; // modified Lentz's algorithm

  // set up for evaluating continued fraction by modified Lentz's method with b0=0.
  // b1 = x+1.-a, c1 = \infty (since c0 = 0), d1 = 1/b1, fr1=b0+a1/b1=a1/b1=1/b1=d1
  double b = x+1.-a, c = 1./FPMIN, d = 1./b, fr = d;
  for(int i = 1; i < ITMAX; i++){
    double aip = -i*(i-a); // a_{i+1}
    b += 2.; d = aip*d+b; // b2, d2 = 1/(a2*d1+b2)
    if(fabs(d) < FPMIN) d = FPMIN;
    c = b+aip/c; // c2 = b2 + a2/c1
    if(fabs(c) < FPMIN) c = FPMIN;
    d = 1./d;
    double delta = c*d;
    fr *= delta; // fr_2 = fr_1*c2*d2
    if(fabs(delta-1.) < EPS)
      return exp(-x+a*log(x)-(gln ? *gln = gammaln(a) : gammaln(a)))*fr; // put factors in front
  } // end for over i
  ETFMsg::Error("ETFMath", "gcf: a too large, and IMAX too small.");
  return 0.; // never gets here
} // end member function gcf

/// fit the track and assign slope and intercept
double ETFMath::IterFit(const int np, const double *z, const double *x,
		const double *r, double &kL, double &bL){
	const int n = pow(2, np);
	double d2, kl, bl, db, dk, q = 0., d2min; // q: goodness of fit
	// dx not unkown, so set to minus deliberately
	vector<double> dx; dx.resize(np); dx.assign(np, -1.);
	LinearFit(z, x, dx.data(), np, bL, kL, db, dk, d2min, q);
	vector<double> zt, xt;
	zt.resize(np); xt.resize(np);
	zt.assign(np, -9999.); xt.assign(np, -9999.);
	const double cosTheta = 1./ sqrt(1.+kL*kL), sinTheta = kL*cosTheta;
	for(int i = n; i--;){
		for(int j = 0; j < np; j++){
			const double R = (2*((i>>j) & 0x1)-1) * r[j];
			zt[j] = z[j] + R*sinTheta;
			xt[j] = x[j] - R*cosTheta;
		} // end for over j
		LinearFit(zt.data(), xt.data(), dx.data(), np, bl, kl, db, dk, d2, q);
		if(d2 < d2min){ kL = kl; bL = bl; d2min = d2; }
	} // end for over i

	return d2min;
} // end member function IterFit

// analytic PID method using DCTaArr
double ETFMath::rho(double k1, double b1, double k2, double b2, double *x2Arr,
	double *zo, double *xo){
	if(fabs(k1 - k2) < 1.e-8){
		ETFMsg::Warn("ETFMath", "rho: |k1-k2| too small: k1: %f, k2: %f", k1, k2);
		k2 = k1 + 1.e-7;
	} // end if

	const double z1 = kzMagIn,  x1 = k1 * z1 + b1; // incident point
	const double z2 = kzMagOut, x2 = k2 * z2 + b2; // exit point

	// x2 solved from equation L1 == L2
	const double x2p = x1+(z2-z1)/(k1+k2)*(k1*k2+sqrt(1.+k1*k1)*sqrt(1.+k2*k2)-1.);
	const double L2p = fabs(((z1-z2)+k1*(x1-x2p))/(k1-k2))*sqrt(1.+k2*k2);

	if(zo && xo){
		zo[0] = (k1*z2-k2*z1+k1*k2*(x2p-x1))/(k1-k2);
		xo[0] = (z1-zo[0])/k1+x1;
	} // end if

#ifdef DEBUG
// XXX: note that L2 is preferred, which agrees with L1p and L2p the better
	const double L1  = fabs(((z1-z2)+k2*(x1-x2 ))/(k1-k2))*sqrt(1.+k1*k1);
	const double L2  = fabs(((z1-z2)+k1*(x1-x2 ))/(k1-k2))*sqrt(1.+k2*k2);
	const double L1p = fabs(((z1-z2)+k2*(x1-x2p))/(k1-k2))*sqrt(1.+k1*k1);

	cout << "\n\nL1: " << L1 << "\tL2: " << L2; // DEBUG
	cout << "\nL1p: " << L1p << "\tL2p: " << L2p; // DEBUG
	cout << "\tx2: " << x2 << "\tx2p: " << x2p << endl; // DEBUG
	getchar(); // DEBUg
#endif

	if(x2Arr){
		x2Arr[0] = x2;
		x2Arr[1] = x2p;
	} // end if
	return L2p;
} // end member function rho

/// \struct tSolve
/// \brief The solution struct for solving particle trajectories in a uniform magnetic field.
struct tSolve{
	double x1, x2, ki, bi, zh, xh, dtheta, rho, zo, xo; // solution struct
};
// solve particle trajectory in uniform magnetic field, with only Mag boundry, exit track
// and target position known; returning the track radius of curvature in the magnetic field
// input unit: mm; output unit: mm
// x=kiz+bi, track before Target
// result: [0-7]: dtheta, rho, (ki, bi), (zo, xo), (x1, x2)
// dtheta: deflection angle; rho, (zo, xo): defines the arc;
// (ki, bi): defines the incident line; (x1, x2): x's of points of into and from magField
// zMagOut->z2; zMagIn->z1; zTa->z0; xTa->x0
bool ETFMath::UniformMagneticSolution(double k1, double b1, double z2, double z1,
		double z0, double x0, double *result){
	const double x2 = k1*z2+b1;
	if(fabs(x2) > 900.) return false; // impossible

	tSolve s[3]; // the equation is of 4-th order, having 4 solutions, one is dropped
	memset(s, 0, sizeof(s));
	for(auto &p : s) p.x2 = x2;

	// transform the equation to a simpler form: x1^3 + b x1^2 + c x1 + d = 0
	const double b = -2.*b1 - x0 - k1*(z0+z1);
	const double c = -x2*x2 + 2.*b1*(x0+x2) + 2.*k1*x2*z0 + z1*(2.*k1*x0-2.*z0+z1) + z2*(2.*z0-z2);
	const double d = -2.*b1*(x0*x2-(z0-z1)*(z1-z2)) +
		x0*(x2*x2-2.*k1*x2*z1+(z1-z2)*(z1-z2)) - k1*(z0-z1)*(x2*x2-z1*z1+z2*z2);

	// let x1 = y - b/3, then the original equation transforms to: y^3 + py + q = 0
	const double p = c - b*b/3.;
	const double q = d + 2.*b*b*b/27. - b*c/3.;

	// solve y^3 + py + q = 0 using Cardano's formula:
	const double qh = q/2., pt = p/3.;
	cdouble Delta(qh*qh + pt*pt*pt, 0.);
	cdouble sqrtDelta = sqrt(Delta);
	cdouble omega(-0.5, Sqrt3()/2.);
	cdouble cqh(qh, 0.);

	cdouble alphap = pow(-cqh + sqrtDelta, 1./3.);
	cdouble alpham = pow(-cqh - sqrtDelta, 1./3.);

	cdouble y[3];
	y[0] = alphap + alpham;
	y[1] = omega*alphap + omega*omega*alpham;
	y[2] = omega*omega*alphap + omega*alpham;

	// output the results and assign its geometrical meaning //
	for(int i = 0; i < 3; i++){
		if(fabs(y[i].imag()) < 1.E-7) s[i].x1 = y[i].real() - b/3.;
		else s[i].x1 = -9999.; // not a real number
		if(fabs(s[i].x1) > 600.) s[i].x1 = -9999.; // out of Mag mouth (+-500)
		if(-9999. == s[i].x1) continue;
		s[i].ki = (s[i].x1 - x0)/(z1 - z0);
		if(fabs(s[i].ki) > 1.45){ s[i].x1 = -9999.; continue; }
		s[i].bi = x0 - s[i].ki*z0;
		s[i].zh = (-b1 + s[i].bi)/(k1 - s[i].ki);
		s[i].xh = k1*s[i].zh + b1;
		s[i].dtheta = atan(k1) - atan(s[i].ki); // anti-clockwise deflection angle
		if(fabs(s[i].dtheta) > 150.*DEGREE()){ s[i].x1 = -9999.; continue; }
		// P1Ph or PhP2 (pdf in misc/tools/math/UniBSolu)
		// if(fabs(s[i].rho) < 500. || fabs(s[i].rho) > 100000.){ s[i].x1 = -9999.; continue; }
		s[i].zo = (k1*(s[i].ki*(s[i].x1 - x2) + z1) - s[i].ki*z2) / (k1 - s[i].ki);
		s[i].xo = (-s[i].ki*s[i].x1 + k1*x2 - z1 + z2) / (k1 - s[i].ki);
		s[i].rho = Distance(s[i].zo, s[i].xo, z2, s[i].x2) * ETFMath::sign(s[i].dtheta);
		// check if (zo, xo)->(z1, x1) == rho and (zo, xo)->(z2, x2) //
		const double r1 = Distance(s[i].zo, s[i].xo, z1, s[i].x1);
		const double r2 = Distance(z2, x2, s[i].zh, s[i].xh)/tan(s[i].dtheta/2.);
		const double r3 = DistanceToLine(s[i].zo, s[i].xo, k1, b1);
		// cout << "rho: " << s[i].rho << " r1: " << r1 <<" r2: " << r2 << " r3: " << r3 << endl; // DEBUG
		// getchar(); // DEBUG
		if(!Within(fabs(r1/s[i].rho), 0.95, 1.05) ||
			 !Within(fabs(r2/s[i].rho), 0.95, 1.05) ||
			 !Within(fabs(r3/s[i].rho), 0.95, 1.05)){
			s[i].x1 = -9999.; continue;
		} // end if


// #define DEBUG_UNI
#ifdef DEBUG_UNI
		const double zo = s[i].zo, xo = s[i].xo; // DEBUG
		const double zh = s[i].zh, xh = s[i].xh; // DEBUG
		const double x1 = s[i].x1; // DEBUG
		const double dd_1 = sqrt(Sum2(z1-zh, x1-xh)); // DEBUG
		const double dd_2 = sqrt(Sum2(z2-zh, x2-xh)); // DEBUG
		const double rho_1 = sqrt(Sum2(z1-zo, x1-xo)); // DEBUG
		const double rho_2 = sqrt(Sum2(z2-zo, x2-xo)); // DEBUG
		cout << "dd_1: " << dd_1 << "\trho_1: " << rho_1 << endl;
		cout << "dd_2: " << dd_2 << "\trho_2: " << rho_2 << endl; // DEBUG
		getchar(); // DEBUG
#endif
	} // end for


	// solution validity check
	int validCnt = 0;
	for(const tSolve &t : s) if(-9999. != t.x1) validCnt++;
	if(1 != validCnt){
		for(int i = 0; i < 6; i++) result[i] = validCnt;
		return false;
	} // end if(1 != validCnt)
	for(const tSolve &t : s) if(-9999. != t.x1){
		// output the result
		result[0] = t.dtheta; result[1] = t.rho;
		result[2] = t.ki; result[3] = t.bi;
		result[4] = t.zo; result[5] = t.xo;
		result[6] = t.x1; result[7] = t.x2;
	} // end for
	return true;
} // end of function UniformMagneticSolution


double ETFMath::BetaGamma(double beta){
	if(beta < 0. || beta >= 1.){
		ETFMsg::Error("ETFMath", "BetaGamma: input beta invalid: %f", beta);
		return -9999.;
	} // end if
	return beta/sqrt(1.-beta*beta);
} // end member function BetaGamma

/// Given a set of data x[0...ndata-1] and y[0...ndata-1] with standard deviations
/// sigy[0...ndata-1]. Fit them to a straight line y(x) = a + bx by minimizing chi2.
/// Returning a and b, and their respective probable uncertainties.siga and sigb, the
/// chi-square chi2 and the goodness-of-fit probability q (that the fit would
/// have chi2 this large or larger under the model given by the fitting). If
/// dy[0] <= 0., then the standard deviations are assumed to be unavailable: q
/// is returned as 1.0 and the normalization of chi2 is to unit standard deviation
/// on all points.
/// Ref. Numerical Recipes in C, p665.
void ETFMath::LinearFit(const double *x, const double *y, const double *dy, int ndata,
    double &a, double &b, double &siga, double &sigb, double &chi2, double &q){
  // stt=\sum_i{1/sigi*(xi-sx/s), sx=\sum_i{xi/sigi^2}, sy=\sum_i{yi/sigi^2}
  double s = 0., sx = 0., sy = 0., stt = 0.;
  bool isSigYKnown = dy[0] > 0.;
  // accumulate sums ...
  if(isSigYKnown){
    s = 0.;
    for(int i = 0; i < ndata; i++){ // ...with weights
      double wt = 1. / (dy[i]*dy[i]); // weights
      s  += wt;
      sx += wt*x[i];
      sy += wt*y[i];
    } // end for
  } // end if
  else{ // ...or without weights
    for(int i = 0; i < ndata; i++){
      sx += x[i];
      sy += y[i];
    } // end for over i
    s = ndata;
  } // end else
  double sxos = sx / s; // sx over s
  // calculate the slope b
  b = 0.;
  if(isSigYKnown){
    for(int i = 0; i < ndata; i++){
      double t = (x[i] - sxos)/dy[i];
      stt += t*t;
      b += t*y[i]/dy[i];
    } // end for
  } // end if
  else{
    for(int i = 0; i < ndata; i++){
      double t = x[i] - sxos;
      stt += t*t;
      b += t*y[i];
    } // end for over i
  } // end else

  // solve for a, b, siga and sigb
  b /= stt;
  a = (sy-sx*b)/s;
  siga = sqrt((1.+sx*sx/(s*stt))/s);
  sigb = sqrt(1./stt);

  // calculte chi2
  chi2 = 0.; q = 0.;
  double nu = ndata - 2; // degree of freedom for chi2 distribution
  if(isSigYKnown){
    for(int i = 0; i < ndata; i++) chi2 += sqr((y[i]-a-b*x[i])/dy[i]);
    q = ETFMath::gammaq(0.5*nu, 0.5*chi2);
  } // end if
  // for unweighted data, evaluate typical sig (which is assumed unit before)
  // using chi2, then update siga and sigb
  else{
    for(int i = 0; i < ndata; i++) chi2 += sqr(y[i]-a-b*x[i]);
    // chi2~nu, chi2*sig^2=\sum{{yi-y(xi)}^2}~nu*sig^2
    double sigData = sqrt(chi2 / nu);
    siga *= sigData;
    sigb *= sigData;
  } // end else

	chi2 /= 1.+b*b; // so that chi2 is sum{d^2}, where d is distance_to_line
} // end member function LinearFit

/////// functions serving 3D tracking //////
/* Mathemathica codes for solving UV->X and UV->Y
tu = -30 Degree;
tv = 30 Degree;
Ry[x_] := {{Cos[x], 0, Sin[x]}, {0, 1, 0}, {-Sin[x], 0, Cos[x]}};
Rz[x_] := {{Cos[x], -Sin[x], 0}, {Sin[x], Cos[x], 0}, {0, 0, 1}};
{xu, yu, zu} = Rz[-tu].Ry[-p].{x, y, z};
{xv, yv, zv} = Rz[-tv].Ry[-p].{x, y, z};
Collect[FullSimplify[Solve[Eliminate[{xu == ku zu + bu, xv == kv zv + bv}, y], x]], z] (*UV\[Rule]X*)
Collect[FullSimplify[Solve[Eliminate[{xu == ku zu + bu, xv == kv zv + bv}, x], y]], z] (*UV\[Rule]Y*)
*/
// Oxuyuzu = Rz[-theta].Ry[-phi].Oxyz: detector is along the zu axis and U wire is parallel to y axis
// Similar rotation has been implemented for V tracking
// U+V->X tranformation: lx: x=kz+b: slope
double ETFMath::kUV_X(double phi, double ku, double kv){
	return ((ku+kv)*cos(phi)+Sqrt3()*sin(phi))/(Sqrt3()*cos(phi)-(ku+kv)*sin(phi));
}
// U+V->X tranformation: lx: x=kz+b: intercept
double ETFMath::bUV_X(double phi, double ku, double kv, double bu, double bv){
	return (bu+bv)/(Sqrt3()*cos(phi)-(ku+kv)*sin(phi));
}
// U+V->Y tranformation: ly: y=kz+b: slope
double ETFMath::kUV_Y(double phi, double ku, double kv){
	return Sqrt3()*(-ku+kv)/(Sqrt3()*cos(phi)-(ku+kv)*sin(phi));
}
// U+V->Y tranformation: ly: y=kz+b: intercept
double ETFMath::bUV_Y(double phi, double ku, double kv, double bu, double bv){
	return (Sqrt3()*(-bu+bv)*cos(phi)+2.*(-bv*ku+bu*kv)*sin(phi))/(Sqrt3()*cos(phi)-(ku+kv)*sin(phi));
}
// X+Y->U tranformation: lu: xu=k zu+bu: slope
double ETFMath::kXY_U(double phi, double k1, double k2){
	return (-k2+Sqrt3()*(k1*cos(phi)-sin(phi)))/(2.*(cos(phi)+k1*sin(phi)));
}
// X+Y->U tranformation: lu: xu=k zu+bu: intercept
double ETFMath::bXY_U(double phi, double k1, double k2, double b1, double b2){
	return 0.5*(-b2+b1*(Sqrt3()+k2*sin(phi))/(cos(phi)+k1*sin(phi)));
}
// X+Y->V tranformation: lv: yv=k zv+bv: slope
double ETFMath::kXY_V(double phi, double k1, double k2){
	return (k2+Sqrt3()*(k1*cos(phi)-sin(phi)))/(2.*(cos(phi)+k1*sin(phi)));
}
// X+Y->V tranformation: lv: yv=k zv+bv: intercept
double ETFMath::bXY_V(double phi, double k1, double k2, double b1, double b2){
	return 0.5*(b2+b1*(Sqrt3()-k2*sin(phi))/(cos(phi)+k1*sin(phi)));
}
